Fix HIP memory, graph, and stream handling. by AWoloszyn · Pull Request #120 · ROCm/hrx-system

AWoloszyn · 2026-06-18T16:59:10Z

Adds HRX-backed HIP pooled allocation support.
Fixes graph capture/launch semantics and graph dependency validation.
Implements stable HIP stream IDs and serialized stream recording.
Tightens stream/context teardown so pending work is synchronized before handles are detached.
Removes stale HIP-local memory-pool helper code after rebasing onto the HRX pool implementation.
Preserves HRX pool buffer ownership for pooled allocations.
Adds overflow/range checks for stream lists, graph memcpy/memset ranges, graph-exec memset sizes, synthetic pointer alignment, and pitched allocations.
Cleans up context failure paths so partially-created contexts and default memory pools are released consistently.

zjgarvey · 2026-06-19T19:20:46Z

+        reachable_nodes, stack);
+  }
+  if (!iree_status_is_ok(status)) {
+    iree_hal_streaming_graph_deinitialize_additional_edge_index(


I think there is a code path where initialize_additional_edge_index calls the deinitialize in an error path and returns an error status, so this additional deinitialize would trigger a double free. Maybe remove the deinit that occurs inside the init func error path since it's handled here.

Good catch.

zjgarvey · 2026-06-19T19:36:56Z

+    HIP_RETURN_ERROR(hipErrorInvalidValue);
+  }
+
+  iree_hal_streaming_graph_node_t* node = NULL;


Why is this a NULL? Is the type enough or something?

Yeah, I guess I just don't understand where the event is going here.

Good catch. Type alone was not enough. Event graph nodes now store and retain the event in the node attrs, and graph exec uses that retained event to
signal/wait. The NULL is now only “no extra trailing node data.”

zjgarvey · 2026-06-22T02:26:09Z

+  }
+
+  iree_hip_per_thread_stream_context = context;
+  iree_hip_per_thread_stream = stream;


I'm a bit confused on how these get cleaned up when the thread exits. I only see the clear function called earlier in this function body, and again during device reset.

Addressed with binding-local TLS. The per-thread stream state has a destructor, so thread exit releases the cached stream and frees the small per-thread
state object. Explicit clears still handle device reset/context replacement.

zjgarvey · 2026-06-22T02:34:21Z

+  return true;
+}
+
+static bool iree_hip_context_invalidate_visible_captures(


This is used in a lot of places, and seems expensive if there isn't any capturing going on. E.g., every hipMalloc needs to snapshot all streams just to see if any of them are capturing?

Is there like a bool we can early return on here if no captures are going on?

zjgarvey · 2026-06-22T02:40:01Z

+    iree_hal_streaming_stream_t* copy_stream =
+        stream ? stream : context->default_stream;
+    iree_hal_streaming_buffer_t* staging = NULL;
+    iree_slim_mutex_lock(&context->mutex);


Why the context mutex and not the copy_stream->mutex?

The context mutex is for the shared pageable H2D staging buffer, not for stream submission. Stream command-buffer mutation is separately protected by
copy_stream->mutex.

zjgarvey · 2026-06-22T02:42:15Z


-    // Record copy command.
+  iree_hal_streaming_stream_t* copy_stream =
+      stream ? stream : context->default_stream;


Does D2H not require any mutexes?

The direct D2H path is a blocking device transfer and does not mutate stream command-buffer state. The fallback command-buffer path does take
copy_stream->mutex while recording.

zjgarvey · 2026-06-22T02:47:17Z

  }

  iree_status_t status = iree_ok_status();
+  if (!stream || stream == hipStreamLegacy || stream == hipStreamPerThread) {


Still trying to understand some of the per-thread stream stuff. Why do we need to synchronize contexts in the per-thread case here?

Agreed, the previous context-wide sync was too broad. hipStreamPerThread should enqueue like an ordinary stream and not synchronize with unrelated streams. I narrowed the NULL/hipStreamLegacy path so it only synchronizes streams that participate in legacy default-stream ordering, excluding hipStreamNonBlocking; per-thread default streams remain blocking streams and only participate in the legacy interop case.

zjgarvey · 2026-06-22T03:04:09Z

-  if (pGraphNode) *pGraphNode = NULL;
-  HIP_RETURN_ERROR(hipErrorNotSupported);
+  IREE_TRACE_ZONE_BEGIN(z0);
+  if (!pGraphNode || !graph || !childGraph) {


Are there really no other checks on the graphs needed here? E.g., can childGraph = graph? Or graph is transitively a child graph of childGraph? I'm not sure what is allowed, or if that's all just up to callers to make sure they don't do some kind of unbounded nesting.

Added validation in the common graph helper. It rejects childGraph == graph, rejects graphs from a different context, and walks nested child graph nodes to reject recursive containment

Retain graph event node state for record and wait nodes and instantiate event-aware graph exec barriers. Validate child graph nodes for context mismatch and recursive containment, while cleaning up retained child graph and event attrs on destroy and clone failure paths. Tighten capture bookkeeping, per-thread stream cleanup, and stream/copy locking around capture adoption and pageable transfer fallbacks.

Use a common thread-local token for capture ownership checks instead of pthread_self. Restore per-thread stream state to existing binding TLS storage and use IREE call-once/mutex primitives for HIP runtime initialization.

Prevent expensive walks where not needed.

AWoloszyn added 6 commits June 18, 2026 11:52

Add HIP pooled allocation backing

1501e39

Fix HIP graph capture and launch semantics

955ebf9

Implement HIP stream handles and serialized recording

7464d4e

Tighten HIP stream and graph cleanup paths

e470912

Harden HIP context cleanup bookkeeping

7258d34

Preserve HRX pool buffer ownership

0faa1d4

AWoloszyn assigned zjgarvey Jun 18, 2026

AWoloszyn requested a review from zjgarvey June 18, 2026 17:52

zjgarvey reviewed Jun 22, 2026

View reviewed changes

AWoloszyn added 5 commits June 23, 2026 09:18

Remove direct pthread use from HIP binding

0beee09

Use a common thread-local token for capture ownership checks instead of pthread_self. Restore per-thread stream state to existing binding TLS storage and use IREE call-once/mutex primitives for HIP runtime initialization.

Update from more comments.

8385798

Narrowed child graph nesting checks.

64642d0

Prevent expensive walks where not needed.

Fix HIP binding CI hygiene

d12ac58

Conversation

AWoloszyn commented Jun 18, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants